import { ChildProcess, IOType, spawn } from 'node:child_process'
import process from 'node:process'
import { Stream } from 'node:stream'
import { ReadBuffer, serializeMessage } from '@modelcontextprotocol/sdk/shared/stdio.js'
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js'

export type StdioServerParameters = {
  /**
   * The executable to run to start the server.
   */
  command: string

  /**
   * Command line arguments to pass to the executable.
   */
  args?: string[]

  /**
   * The environment to use when spawning the process.
   *
   * If not specified, the result of getDefaultEnvironment() will be used.
   */
  env?: Record<string, string>

  /**
   * How to handle stderr of the child process. This matches the semantics of Node's `child_process.spawn`.
   *
   * The default is "inherit", meaning messages to stderr will be printed to the parent process's stderr.
   */
  stderr?: IOType | Stream | number

  /**
   * The working directory to use when spawning the process.
   *
   * If not specified, the current working directory will be inherited.
   */
  cwd?: string
}

/**
 * Environment variables to inherit by default, if an environment is not explicitly given.
 */
export const DEFAULT_INHERITED_ENV_VARS =
  process.platform === 'win32'
    ? [
        'APPDATA',
        'HOMEDRIVE',
        'HOMEPATH',
        'LOCALAPPDATA',
        'PATH',
        'PROCESSOR_ARCHITECTURE',
        'SYSTEMDRIVE',
        'SYSTEMROOT',
        'TEMP',
        'USERNAME',
        'USERPROFILE'
      ]
    : /* list inspired by the default env inheritance of sudo */
      ['HOME', 'LOGNAME', 'PATH', 'SHELL', 'TERM', 'USER']

/**
 * Returns a default environment object including only environment variables deemed safe to inherit.
 */
export function getDefaultEnvironment(): Record<string, string> {
  const env: Record<string, string> = {}

  for (const key of DEFAULT_INHERITED_ENV_VARS) {
    const value = process.env[key]
    if (value === undefined) {
      continue
    }

    if (value.startsWith('()')) {
      // Skip functions, which are a security risk.
      continue
    }

    env[key] = value
  }

  return env
}

/**
 * Client transport for stdio: this will connect to a server by spawning a process and communicating with it over stdin/stdout.
 *
 * This transport is only available in Node.js environments.
 */
export class StdioClientTransport implements Transport {
  private _process?: ChildProcess
  private _abortController: AbortController = new AbortController()
  private _readBuffer: ReadBuffer = new ReadBuffer()
  private _serverParams: StdioServerParameters

  onclose?: () => void
  onerror?: (error: Error) => void
  onmessage?: (message: JSONRPCMessage) => void

  constructor(server: StdioServerParameters) {
    this._serverParams = server
  }

  /**
   * Starts the server process and prepares to communicate with it.
   */
  async start(): Promise<void> {
    if (this._process) {
      throw new Error(
        'StdioClientTransport already started! If using Client class, note that connect() calls start() automatically.'
      )
    }

    return new Promise((resolve, reject) => {
      this._process = spawn(this._serverParams.command, this._serverParams.args ?? [], {
        env: this._serverParams.env ?? getDefaultEnvironment(),
        stdio: ['pipe', 'pipe', this._serverParams.stderr ?? 'inherit'],
        shell: process.platform === 'win32' ? true : false,
        signal: this._abortController.signal,
        windowsHide: process.platform === 'win32' && isElectron(),
        cwd: this._serverParams.cwd
      })

      this._process.on('error', (error) => {
        if (error.name === 'AbortError') {
          // Expected when close() is called.
          this.onclose?.()
          return
        }

        reject(error)
        this.onerror?.(error)
      })

      this._process.on('spawn', () => {
        resolve()
      })

      this._process.on('close', (_code) => {
        this._process = undefined
        this.onclose?.()
      })

      this._process.stdin?.on('error', (error) => {
        this.onerror?.(error)
      })

      this._process.stdout?.on('data', (chunk) => {
        this._readBuffer.append(chunk)
        this.processReadBuffer()
      })

      this._process.stdout?.on('error', (error) => {
        this.onerror?.(error)
      })
    })
  }

  /**
   * The stderr stream of the child process, if `StdioServerParameters.stderr` was set to "pipe" or "overlapped".
   *
   * This is only available after the process has been started.
   */
  get stderr(): Stream | null {
    return this._process?.stderr ?? null
  }

  private processReadBuffer() {
    while (true) {
      try {
        const message = this._readBuffer.readMessage()
        if (message === null) {
          break
        }

        this.onmessage?.(message)
      } catch (error) {
        this.onerror?.(error as Error)
      }
    }
  }

  async close(): Promise<void> {
    this._abortController.abort()
    this._process = undefined
    this._readBuffer.clear()
  }

  send(message: JSONRPCMessage): Promise<void> {
    return new Promise((resolve) => {
      if (!this._process?.stdin) {
        throw new Error('Not connected')
      }

      const json = serializeMessage(message)
      if (this._process.stdin.write(json)) {
        resolve()
      } else {
        this._process.stdin.once('drain', resolve)
      }
    })
  }
}

function isElectron() {
  return 'type' in process
}
